[XNNPACK] Lower upsample_nearest2d x2 as concatenate4 + depth_to_space_2d - #24113

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24113

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• PRs based before #202327 fail all Linux builds: can't find build-external-packages action

❌ 2 New Failures

As of commit b2eb31c with merge base 422cc39 ():

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@pytorchbot label "release notes: xnnpack"

sorry, can we split this into two PRs? to keep the changes atomic per pr instead of squishing them together. the upsample changes are additive and lower risk.

that being said, a lower-rank operand with symbolic dims breaks lowering with these changes. Before, a dynamic-shape model with conv_out + y, where y is [H, W] or [C, H, W] with dynamic H/W, now fails in preprocess. On main it lowered fine by falling back to NCHW.

sorry, can we split this into two PRs? to keep the changes atomic per pr instead of squishing them together. the upsample changes are additive and lower risk.

that being said, a lower-rank operand with symbolic dims breaks lowering with these changes. Before, a dynamic-shape model with conv_out + y, where y is [H, W] or [C, H, W] with dynamic H/W, now fails in preprocess. On main it lowered fine by falling back to NCHW.

I will look into that tomorrow

Done, split in two: this PR now only lowers upsample_nearest2d, the ChannelsLastTaggedReshapePass change moved to #24183.

Thanks for the dynamic-shape report, reproduced it: unsqueeze_to_4d put the symbolic dims into the view_copy shape, which XNNStaticReshape rejects. In #24183 an operand with symbolic dims keeps the NCHW path as on main, with a test for conv_out + y with y of shape [H, W] and dynamic H/W that also runs at a second input size.

a few things:

1. A quantized model can fail to load. xnn_define_depth_to_space_2d rejects the node unless its input and output have the same quantization parameters. The visitor uses the input's parameters for the input and separate output parameters for the output. The partitioner should check that the two match, or set up the quantization so they always do.
2. Wrong results when scale_factors is a kwarg as the check that the scale is exactly 2.0 only reads node.args[2]. Take a 200×200 input with scale_factors=[2.004, 2.004] passed as a kwarg: the output is still 400×400, so the size check passes and the scale check is skipped.
3. The qs8 test can't fail (test_upsample_nearest2d.py:61). check_not(["torch.ops.quantized_decomposed"]) runs on the edge graph, where the q/dq ops have different names, so it never matches.
4. The annotator covers more than the partitioner accepts so it quantizes every upsample_nearest2d.vec

Thanks, all four addressed in b2eb31c.

1. Mismatched quantization parameters: handled in the visitor rather than rejected. When the output is quantized differently from the input, depth_to_space_2d writes into an intermediate carrying the input's parameters and an XNNConvert requantizes it to the output (XNNPACK supports qint8 to qint8 convert). I first tried rejecting the node in the partitioner, but since the surrounding q/dq get delegated on their own, the not-decompose sanity check in to_edge_transform_and_lower re-runs the constraint on the partitioned graph, where the q/dq context is gone, and raises. Test: test_qs8_upsample_nearest2d_requantized_output, which fails at load without the convert.
2. scale_factors as a kwarg: the check reads node.kwargs when the positional argument is absent. torch.export and to_edge normalize it to positional, so I could not reproduce it end to end; the test rewrites the exported node to a kwarg and asserts on check_constraints directly.
3. The qs8 test: the check_not on the edge graph was indeed vacuous, so it is gone. The test now checks on the exported graph that the model is quantized and relies on the single-delegate and no-edge-op checks after lowering. test_add.py and test_cat.py carry the same vacuous check_not; left untouched here.
4. Annotator scope: the x2 constraint moved to why_not_x2_nearest_upsample in utils.py, used by both UpsampleNearest2dConfig and the annotator, so the quantizer only annotates what the partitioner lowers. Test: test_qs8_upsample_nearest2d_unsupported_scale_not_annotated for x3 and 2.004.